import logging
from pathlib import PurePath
from typing import Callable, Collection, Set, Type

from common import OperatingSystem
from common.agent_events import AgentEventTag
from common.credentials import Credentials, NTHash
from common.tags import (
    ALTERNATE_AUTHENTICATION_MATERIAL_T1550_TAG,
    BRUTE_FORCE_T1110_TAG,
    INGRESS_TOOL_TRANSFER_T1105_TAG,
    REMOTE_SERVICES_T1021_TAG,
)
from common.types import NetworkPort
from infection_monkey.exploit.tools import (
    IRemoteAccessClient,
    RemoteAuthenticationError,
    RemoteCommandExecutionError,
    RemoteFileCopyError,
)
from infection_monkey.i_puppet import TargetHost

from .rdp_client import RDPClient
from .rdp_options import RDPOptions

logger = logging.getLogger(__name__)


LOGIN_TAGS = {
    REMOTE_SERVICES_T1021_TAG,
    BRUTE_FORCE_T1110_TAG,
}
COPY_FILE_TAGS = {
    REMOTE_SERVICES_T1021_TAG,
    INGRESS_TOOL_TRANSFER_T1105_TAG,
}
EXECUTION_TAGS = {
    REMOTE_SERVICES_T1021_TAG,
}

RDP_PORTS = [NetworkPort(3389)]


class RDPRemoteAccessClient(IRemoteAccessClient):
    def __init__(
        self,
        host: TargetHost,
        options: RDPOptions,
        command_builder: Callable[[PurePath], str],
        rdp_client: RDPClient,
    ):
        self._host = host
        self._options = options
        self._build_command = command_builder
        self._client = rdp_client

    def login(self, credentials: Credentials, tags: Set[AgentEventTag]):
        tags.update(LOGIN_TAGS)
        try:
            if isinstance(credentials.secret, NTHash):
                tags.update({ALTERNATE_AUTHENTICATION_MATERIAL_T1550_TAG})

            self._client.connect(
                self._host, credentials, RDP_PORTS[0], self._options.rdp_connect_timeout
            )
        except Exception as err:
            raise RemoteAuthenticationError("Failed to login to RDP server") from err

    def _raise_if_not_authenticated(self, error_type: Type[Exception]):
        if not self._client.connected():
            raise error_type(
                "This operation cannot be performed until authentication is successful"
            )

    def get_os(self) -> OperatingSystem:
        return OperatingSystem.WINDOWS

    def execute_agent(self, agent_binary_path: PurePath, tags: Set[AgentEventTag]):
        self._raise_if_not_authenticated(RemoteCommandExecutionError)
        tags.update(EXECUTION_TAGS)
        try:
            command = self._build_command(remote_agent_binary_destination_path=agent_binary_path)
            self._client.execute_command(command)
        except Exception as err:
            raise RemoteCommandExecutionError("Failed to execute agent") from err

    def copy_file(self, file: bytes, destination_path: PurePath, tags: Set[AgentEventTag]):
        self._raise_if_not_authenticated(RemoteFileCopyError)

        logger.debug(
            f"Trying to copy monkey file to [{destination_path}] on victim {self._host.ip}"
        )

        tags.update(COPY_FILE_TAGS)
        try:
            self._client.copy_file(
                file,
                destination_path,
                self._options.login_timeout,
                self._options.agent_binary_upload_timeout,
            )
        except Exception as err:
            raise RemoteFileCopyError("Failed to copy file to RDP server") from err

    def get_writable_paths(self) -> Collection[PurePath]:
        return []
